#coding=utf-8
u"""
一个多烟团模型，完全采用numpy和numexpr实现，目标是为了给GPU运算做基准测试
numpy和numexpr都采用了MKL加速
空间复杂度O(n)=GRID_SIZE^2*puff*O(exp)
时间复杂度o(n)=2*O(n)*timepoints
平均一个时刻的一个烟团在采用numexpr加速的情况下需要
GRID-SIZE       TIME
4096*4096       3.4s
2048*2048       1.8s  -> current choice
1024*1024       0.3s
"""

## Here is some note:
## Walking-length is a tuple (xy, xz). At beginning, xy equals xz until the stability changed so new walking-length
## are calculated.
## The position of smoke is a tuple (x, y, z). Remember x, y and z are all float numbers, not integers.
## Met-condition is a tuple (u,v,z,stab)
## Z-speed is ignored in this version of puff-model.
## Time interval of met-condition matrix is HARD-CODED as ten-minutes, but the time interval of source-releasing defined
## in TIMESTEP (in seconds). Also, the TIMESTEP are time interval to track each smoke.

## FOR DEBUG USE:
if __debug__:
    import matplotlib.pyplot as plt

import os
os.environ["CUDA_PROFILE"] = "1"
os.environ["CUDA_DEVICE"] = "0"

import numpy
import numexpr
import math
import datetime; from datetime import *

from global_settings import *
import sample_gen
try:
    import testing_class
    from testing_class import run_gpu_conc
except:
    cuda_disabled = True
else:
    cuda_disabled = False

#Set up numexpr
numexpr.set_vml_accuracy_mode("high")

#建立坐标网格
HG = GRID_SIZE/2
GridIdxY, GridIdxX = numpy.mgrid[HG-1:-HG-1:-1, -HG:HG]
CentralPositionMatrixX = numexpr.evaluate("GridIdxX*GRID_INTERVAL+0.5*GRID_INTERVAL")
CentralPositionMatrixY = numexpr.evaluate("GridIdxY*GRID_INTERVAL-0.5*GRID_INTERVAL")

#读取输入数据
testdata = sample_gen.MetPreProcessor(mode=2)
print "Reading source data"
ReleaseQ = testdata.GenerateTestSource()
print "Reading met data"
MetField = testdata.GenerateField()
print "Reading met sequence"
MetSeq = testdata.Generate_MetSeq()

smoke_list = []
result_list = []


class model_wrf_sam:

    def __init__(release):
        self.release = release
        self.met_field = met_field
        self.met_seq = met_seq

    def runbenchmark(self, height_in=1):
        tick = 0
        height = height_in
        #Total Estimation Time is DURATION seconds
        while tick * TIMESTEP < DURATION:
            #判断烟团释放过程是否结束
            if tick < len(self.release):
                #表示此刻有无烟团释放
                if self.release[tick] > 0:
                    #加入烟团list
                    smoke_list.append(smoke_def(self.release[tick], met_field, tick))
            #追踪当前每个烟团的位置，并且计算浓度场
            empty_conc_matrix = numpy.empty_like(CentralPositionMatrixX)
            for smoke in smoke_list:
                if not smoke.invalid:
                    smoke.walk()
                #只有当需要输出的时刻才计算浓度场
                if not smoke.invalid and tick % OUTSTEP == 0:
                    #Prepare some variables used in numexpr.evaluate
                    X, Y, Z = smoke.pos
                    Stability = smoke.met[3]
                    mass = smoke.mass
                    #NOTE: walkinglength should be calculated when met condition changes
                    DiffC = smoke.diffusion_coefficents(Stability, smoke.walkinglength)
                    x, y, z = DiffC
                    PI = math.pi
                    current_smoke_conc = numexpr.evaluate("mass/((2*PI)**1.5*x*y*z)*exp(-0.5*((CentralPositionMatrixX-X)/x)**2)*exp(-0.5*((CentralPositionMatrixY-Y)/y)**2)*(exp(-0.5*((Z-height)/z)**2) + exp(-0.5*((Z+height)/z)**2))")
                    empty_conc_matrix = numexpr.evaluate("empty_conc_matrix + current_smoke_conc")

            result_list.append(empty_conc_matrix)
            #时钟+1
            tick += 1
            print "current time-tick is %d" % tick
            ##FOR DEBUG USE
            if __debug__ and tick == 80:
                print numpy.max(empty_conc_matrix)
                plot_debug(empty_conc_matrix)
                return -1
        return 0

    def run_gpu_benchmark(self, height_in=1):
        tick = 0
        #做一点赋值，方便调试
        gridw = GRID_INTERVAL
        GRID_WIDTH = GRID_SIZE
        height = height_in
        #Total Estimation Time is DURATION seconds
        while tick * TIMESTEP < DURATION:
            #判断烟团释放过程是否结束
            if tick < len(self.release):
                #表示此刻有无烟团释放
                if self.release[tick] > 0:
                    #加入烟团list
                    smoke_list.append(smoke_def(self.release[tick], met_field, tick))
            #追踪当前每个烟团的位置，并且计算浓度场
            walkinglength = []; mass = []; stab = []; pos = [];
            for smoke in smoke_list:
                if not smoke.invalid:
                    smoke.walk()
                #只有当需要输出的时刻才计算浓度场
                if not smoke.invalid and tick % OUTSTEP == 0:
                    pos += list(smoke.pos)
                    walkinglength += list(smoke.walkinglength)
                    stab.append(int(smoke.met[3]))
                    mass.append(smoke.mass)
            if len(mass) == 0:
                print "no valid mass at tick %d" % tick
                tick += 1
                continue
            dest = run_gpu_conc(pos=numpy.array(pos).astype(numpy.float32), walkinglength=numpy.array(walkinglength).astype(numpy.float32), mass=numpy.array(mass).astype(numpy.float32),
                                        height=height, GRID_WIDTH=GRID_WIDTH, stab=numpy.array(stab).astype(numpy.int32), gridw=gridw, smoke_count=len(mass))
            empty_conc_matrix = dest.reshape((GRID_SIZE,GRID_SIZE), order='C')

            result_list.append(empty_conc_matrix)
            #时钟+1
            tick += 1
            print "current time-tick is %d" % tick
            ##debug code
            if __debug__ and tick == 80:
                print numpy.max(empty_conc_matrix)
                plot_debug(empty_conc_matrix)
                return -1
        return 0

#debug use
def plot_debug(result):
    plt.contour(result)
    plt.contourf(result)
    plt.show()

#测试
startcpu = datetime.now()
if not cuda_disabled:
    print "We found a cuda card and pycuda initialized."
    result = run_gpu_benchmark()
else:
    print "We found no cuda card. We'll use cpu only."
    resutl = runbenchmark()
print datetime.now() - startcpu

#result = get_grid_concentration(pos=Pos, walkinglength=WalkingLength, mass=ReleaseQ[0])






